import pytest
from matplotlib._api.__init__ import select_matching_signature

# unit tests

# ---- BASIC TEST CASES ----


def test_single_function_exact_signature():
    # Test with one function that matches the arguments exactly
    def f(a, b):
        return a + b

    codeflash_output = select_matching_signature([f], 2, 3)
    result = codeflash_output  # 1.31μs -> 1.02μs (27.9% faster)


def test_multiple_functions_first_matches():
    # First function matches, second would also match but should not be called
    called = []

    def f1(a):
        called.append(1)
        return a * 2

    def f2(a):
        called.append(2)
        return a * 3

    codeflash_output = select_matching_signature([f1, f2], 4)
    result = codeflash_output  # 1.60μs -> 1.21μs (32.4% faster)


def test_multiple_functions_second_matches():
    # First function fails, second matches
    def f1(a, b):
        raise Exception("Should not be called if signature doesn't match")

    def f2(a):
        return a * 2

    codeflash_output = select_matching_signature([f1, f2], 3)
    result = codeflash_output  # 4.13μs -> 3.26μs (26.7% faster)


def test_kwargs_matching():
    # Test with keyword arguments
    def f(a, b=10):
        return a + b

    codeflash_output = select_matching_signature([f], a=5, b=7)
    result = codeflash_output  # 1.79μs -> 1.49μs (19.8% faster)


def test_no_matching_signature_raises_typeerror():
    # None of the functions match the signature, should raise TypeError from the last
    def f1(a):
        pass

    def f2(b):
        pass

    with pytest.raises(TypeError):
        select_matching_signature([f1, f2], 1, 2)  # 5.15μs -> 4.71μs (9.32% faster)


def test_function_raises_typeerror_internally():
    # Only TypeError due to signature should be caught, not TypeError inside the function body
    def f1(a):
        raise TypeError("internal error")

    with pytest.raises(TypeError) as excinfo:
        select_matching_signature([f1], 1)  # 2.08μs -> 1.83μs (13.8% faster)


def test_function_with_varargs():
    # Test function with *args signature
    def f1(*args):
        return sum(args)

    codeflash_output = select_matching_signature([f1], 1, 2, 3)
    result = codeflash_output  # 1.90μs -> 1.46μs (30.4% faster)


def test_function_with_kwargs():
    # Test function with **kwargs signature
    def f1(**kwargs):
        return kwargs["x"] * 2

    codeflash_output = select_matching_signature([f1], x=4)
    result = codeflash_output  # 1.90μs -> 1.58μs (20.4% faster)


def test_function_with_default_values():
    # Function with default values for arguments
    def f1(a, b=5):
        return a * b

    codeflash_output = select_matching_signature([f1], 3)
    result = codeflash_output  # 1.46μs -> 1.05μs (38.6% faster)


def test_function_with_positional_and_keyword_args():
    # Function with both positional and keyword arguments
    def f1(a, b, c=0):
        return a + b + c

    codeflash_output = select_matching_signature([f1], 1, 2, c=3)
    result = codeflash_output  # 1.83μs -> 1.51μs (21.3% faster)


def test_function_with_missing_required_kwarg():
    # Should raise TypeError if required kwarg is missing
    def f1(a, b):
        return a + b

    with pytest.raises(TypeError):
        select_matching_signature([f1], a=1)  # 4.02μs -> 3.58μs (12.3% faster)


def test_function_with_extra_kwargs():
    # Should raise TypeError if extra kwarg is passed
    def f1(a):
        return a

    with pytest.raises(TypeError):
        select_matching_signature([f1], a=1, b=2)  # 3.02μs -> 2.67μs (13.0% faster)


def test_function_with_extra_args():
    # Should raise TypeError if extra positional arg is passed
    def f1(a):
        return a

    with pytest.raises(TypeError):
        select_matching_signature([f1], 1, 2)  # 4.00μs -> 3.66μs (9.22% faster)


def test_function_with_signature_and_varargs():
    # Should prefer signature match over *args
    def f1(a, b):
        return a + b

    def f2(*args):
        return sum(args)

    codeflash_output = select_matching_signature([f1, f2], 2, 3)
    result = codeflash_output  # 1.37μs -> 1.07μs (27.7% faster)


def test_function_with_signature_and_kwargs():
    # Should prefer signature match over **kwargs
    def f1(a, b):
        return a * b

    def f2(**kwargs):
        return kwargs["a"] + kwargs["b"]

    codeflash_output = select_matching_signature([f1, f2], a=2, b=3)
    result = codeflash_output  # 1.87μs -> 1.52μs (22.7% faster)


def test_function_with_signature_and_varargs_kwargs():
    # Should prefer signature match over *args, **kwargs
    def f1(a, b):
        return a - b

    def f2(*args, **kwargs):
        return sum(args) + sum(kwargs.values())

    codeflash_output = select_matching_signature([f1, f2], 7, 2)
    result = codeflash_output  # 1.37μs -> 1.03μs (33.2% faster)


def test_function_with_signature_and_failing_varargs():
    # If signature fails, *args fallback should be used
    def f1(a, b):
        raise TypeError("internal error")

    def f2(*args):
        return len(args)

    codeflash_output = select_matching_signature([f1, f2], 1, 2, 3)
    result = codeflash_output  # 4.29μs -> 3.68μs (16.4% faster)


def test_function_with_signature_and_failing_kwargs():
    # If signature fails, **kwargs fallback should be used
    def f1(a, b):
        raise TypeError("internal error")

    def f2(**kwargs):
        return len(kwargs)

    codeflash_output = select_matching_signature([f1, f2], a=1, b=2, c=3)
    result = codeflash_output  # 3.86μs -> 3.11μs (24.0% faster)


def test_function_with_signature_and_failing_all():
    # All functions fail, should raise TypeError from the last
    def f1(a):
        raise TypeError("fail 1")

    def f2(b):
        raise TypeError("fail 2")

    with pytest.raises(TypeError) as excinfo:
        select_matching_signature([f1, f2], 1)  # 2.77μs -> 2.37μs (16.9% faster)


def test_function_with_signature_and_non_typeerror():
    # If a function raises a non-TypeError, it should propagate
    def f1(a):
        raise ValueError("not a type error")

    with pytest.raises(ValueError):
        select_matching_signature([f1], 1)  # 1.89μs -> 1.61μs (17.2% faster)


# ---- LARGE SCALE TEST CASES ----


def test_large_number_of_functions_first_matches():
    # Large list of functions, first matches
    def make_func(i):
        return lambda x: x + i

    funcs = [make_func(i) for i in range(1000)]
    codeflash_output = select_matching_signature(funcs, 5)
    result = codeflash_output  # 1.63μs -> 1.12μs (45.8% faster)


def test_large_number_of_functions_last_matches():
    # Large list of functions, last matches
    def fail(*args, **kwargs):
        raise TypeError("fail")

    def succeed(x):
        return x * 2

    funcs = [fail] * 999 + [succeed]
    codeflash_output = select_matching_signature(funcs, 7)
    result = codeflash_output  # 234μs -> 178μs (31.6% faster)


def test_large_number_of_functions_none_match():
    # Large list of functions, none match, should raise TypeError
    def fail1(x):
        raise TypeError("fail1")

    def fail2(x):
        raise TypeError("fail2")

    funcs = [fail1] * 500 + [fail2] * 500
    with pytest.raises(TypeError) as excinfo:
        select_matching_signature(funcs, 1, 2)  # 583μs -> 543μs (7.33% faster)


def test_large_number_of_functions_some_raise_other_error():
    # Some functions raise ValueError, only TypeError should be caught
    def fail_type(x):
        raise TypeError("fail type")

    def fail_value(x):
        raise ValueError("fail value")

    def succeed(x):
        return x + 1

    funcs = [fail_type] * 300 + [fail_value] * 300 + [succeed]
    with pytest.raises(ValueError):
        select_matching_signature(funcs, 1)  # 63.4μs -> 48.9μs (29.7% faster)


def test_large_number_of_functions_with_various_signatures():
    # Mix of signatures, only one matches
    def f1(a, b):
        return a + b

    def f2(*args):
        return sum(args)

    def f3(**kwargs):
        return sum(kwargs.values())

    funcs = [f1] * 333 + [f2] * 333 + [f3] * 334
    codeflash_output = select_matching_signature(funcs, 1, 2)
    result = codeflash_output  # 1.35μs -> 1.11μs (21.5% faster)
    codeflash_output = select_matching_signature(funcs, 1, 2, 3)
    result2 = codeflash_output  # 196μs -> 180μs (8.93% faster)
    codeflash_output = select_matching_signature(funcs, a=1, b=2)
    result3 = codeflash_output  # 1.20μs -> 1.17μs (2.57% faster)


# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

import pytest
from matplotlib._api.__init__ import select_matching_signature

# unit tests

# ---------------- BASIC FUNCTIONALITY TESTS ----------------


def test_single_function_matching_args():
    # Only one function, correct signature
    def f(a, b):
        return a + b

    codeflash_output = select_matching_signature(
        [f], 2, 3
    )  # 1.34μs -> 1.02μs (31.4% faster)


def test_multiple_functions_first_matches():
    # First function matches, others would also match but should not be called
    called = []

    def f1(x):
        called.append(1)
        return x + 1

    def f2(x):
        called.append(2)
        return x + 2

    codeflash_output = select_matching_signature([f1, f2], 10)
    result = codeflash_output  # 1.55μs -> 1.19μs (30.1% faster)


def test_multiple_functions_second_matches():
    # First function fails, second matches
    def f1(x, y):
        raise AssertionError("Should not be called")

    def f2(x):
        return x * 2

    codeflash_output = select_matching_signature(
        [f1, f2], 4
    )  # 4.02μs -> 3.18μs (26.2% faster)


def test_kwargs_passed():
    # Functions with kwargs
    def f1(a):
        return a

    def f2(a, b=0):
        return a + b

    codeflash_output = select_matching_signature(
        [f1, f2], 5, b=7
    )  # 3.34μs -> 2.65μs (26.0% faster)


def test_function_with_varargs():
    # Function with *args
    def f1(*args):
        return sum(args)

    codeflash_output = select_matching_signature(
        [f1], 1, 2, 3, 4
    )  # 1.91μs -> 1.56μs (22.4% faster)


def test_function_with_varkwargs():
    # Function with **kwargs
    def f1(**kwargs):
        return kwargs["x"] + kwargs["y"]

    codeflash_output = select_matching_signature(
        [f1], x=2, y=3
    )  # 2.07μs -> 1.69μs (22.5% faster)


def test_function_with_args_and_kwargs():
    def f1(a, **kwargs):
        return a + kwargs.get("b", 0)

    codeflash_output = select_matching_signature(
        [f1], 3, b=4
    )  # 1.88μs -> 1.67μs (12.6% faster)


# ---------------- EDGE CASES ----------------


def test_no_matching_signature():
    # None of the functions match, should raise TypeError
    def f1(a):
        return a

    def f2(b):
        return b

    with pytest.raises(TypeError):
        select_matching_signature([f1, f2], 1, 2, 3)  # 5.12μs -> 4.58μs (11.7% faster)


def test_all_functions_raise_typeerror():
    # All functions raise TypeError due to signature mismatch
    def f1(a):
        return a

    def f2(b):
        return b

    with pytest.raises(TypeError):
        select_matching_signature([f1, f2], 1, 2)  # 4.61μs -> 4.50μs (2.47% faster)


def test_function_raises_non_typeerror():
    # If a function raises a non-TypeError, it should propagate
    def f1(a):
        raise ValueError("Custom error")

    def f2(a):
        return a

    with pytest.raises(ValueError):
        select_matching_signature([f1, f2], 1)  # 1.97μs -> 1.60μs (22.8% faster)


def test_function_returns_none():
    # Matching function returns None
    def f1(a):
        return None

    codeflash_output = select_matching_signature(
        [f1], 1
    )  # 1.33μs -> 983ns (35.8% faster)


def test_function_with_default_args():
    # Function with default arguments
    def f1(a, b=10):
        return a + b

    codeflash_output = select_matching_signature(
        [f1], 5
    )  # 1.38μs -> 1.02μs (34.6% faster)


def test_function_with_positional_and_keyword_only():
    def f1(a, *, b):
        return a * b

    codeflash_output = select_matching_signature(
        [f1], 3, b=4
    )  # 1.81μs -> 1.45μs (24.3% faster)


def test_function_with_no_args():
    def f1():
        return 42

    codeflash_output = select_matching_signature([f1])  # 1.29μs -> 928ns (38.8% faster)


def test_function_with_args_and_kwargs_and_extra_kwargs():
    def f1(a, **kwargs):
        return a + kwargs.get("b", 0)

    codeflash_output = select_matching_signature(
        [f1], 2, b=3, c=100
    )  # 2.19μs -> 1.82μs (20.8% faster)


def test_function_with_args_and_kwargs_and_missing_kwarg():
    def f1(a, b):
        return a + b

    def f2(a):
        return a

    codeflash_output = select_matching_signature(
        [f1, f2], 7
    )  # 3.89μs -> 3.17μs (22.6% faster)


def test_function_with_args_and_kwargs_and_extra_args():
    def f1(a):
        return a

    def f2(a, b):
        return a + b

    codeflash_output = select_matching_signature(
        [f1, f2], 1, 2
    )  # 4.20μs -> 3.53μs (19.1% faster)


def test_typeerror_message_propagation():
    # The TypeError message from the last function is what is raised
    def f1(a):
        return a

    def f2():
        raise TypeError("custom signature error")

    with pytest.raises(TypeError) as excinfo:
        select_matching_signature([f1, f2], 1, 2)  # 5.25μs -> 4.73μs (11.0% faster)


# ---------------- LARGE SCALE TESTS ----------------


def test_large_number_of_functions_first_matches():
    # 1000 functions, first matches
    def make_func(i):
        def f(x):
            return x + i

        return f

    funcs = [make_func(i) for i in range(1000)]
    codeflash_output = select_matching_signature(
        funcs, 10
    )  # 1.59μs -> 1.05μs (50.7% faster)


def test_large_number_of_functions_last_matches():
    # 999 functions fail, last matches
    def make_func(i):
        def f(x):
            raise TypeError()

        return f

    def last_func(x):
        return x * 2

    funcs = [make_func(i) for i in range(999)] + [last_func]
    codeflash_output = select_matching_signature(
        funcs, 7
    )  # 195μs -> 149μs (30.1% faster)


def test_large_number_of_functions_all_fail():
    # All 1000 functions fail
    def make_func(i):
        def f(x):
            raise TypeError(f"fail {i}")

        return f

    funcs = [make_func(i) for i in range(1000)]
    with pytest.raises(TypeError) as excinfo:
        select_matching_signature(funcs, 1)  # 297μs -> 236μs (25.6% faster)


def test_large_number_of_functions_middle_matches():
    # Match in the middle
    def fail(x):
        raise TypeError()

    def match(x):
        return x * 10

    funcs = [fail] * 500 + [match] + [fail] * 499
    codeflash_output = select_matching_signature(
        funcs, 3
    )  # 99.8μs -> 75.7μs (31.8% faster)


def test_large_number_of_functions_with_various_signatures():
    # Many functions with different signatures, only one matches
    def make_func(i):
        if i == 500:
            return lambda x, y=0: x + y
        else:
            return lambda x: x

    funcs = [make_func(i) for i in range(1000)]
    codeflash_output = select_matching_signature(
        funcs, 7, y=3
    )  # 169μs -> 139μs (22.0% faster)


# ---------------- AMBIGUITY TESTS ----------------


def test_ambiguous_signatures_returns_first():
    # Both could match, but only first is called
    def f1(a, b=0):
        return a + b

    def f2(a, b=0):
        return a * b

    codeflash_output = select_matching_signature(
        [f1, f2], 2, b=3
    )  # 1.75μs -> 1.46μs (20.0% faster)


def test_ambiguous_signatures_second_matches():
    # First fails, second matches
    def f1(a):
        raise TypeError()

    def f2(a, b=0):
        return a * b

    codeflash_output = select_matching_signature(
        [f1, f2], 2, b=4
    )  # 3.35μs -> 2.66μs (25.7% faster)


# ---------------- MISCELLANEOUS TESTS ----------------


def test_function_with_args_and_kwargs_and_weird_args():
    # Function with *args and **kwargs, unexpected extra args
    def f1(*args, **kwargs):
        return sum(args) + sum(kwargs.values())

    codeflash_output = select_matching_signature(
        [f1], 1, 2, 3, x=4, y=5
    )  # 2.97μs -> 2.68μs (10.9% faster)


def test_function_with_keyword_only_args():
    def f1(*, x, y):
        return x * y

    codeflash_output = select_matching_signature(
        [f1], x=2, y=3
    )  # 1.98μs -> 1.62μs (22.4% faster)


def test_function_with_positional_only_args():
    def f1(a, /, b):
        return a + b

    codeflash_output = select_matching_signature(
        [f1], 3, 4
    )  # 1.37μs -> 1.06μs (29.1% faster)


def test_function_with_mixed_args():
    def f1(a, /, b, *, c):
        return a + b + c

    codeflash_output = select_matching_signature(
        [f1], 1, 2, c=3
    )  # 1.81μs -> 1.49μs (21.5% faster)


def test_function_with_varargs_and_varkwargs():
    def f1(*args, **kwargs):
        return len(args) + len(kwargs)

    codeflash_output = select_matching_signature(
        [f1], 1, 2, a=3, b=4
    )  # 2.24μs -> 1.84μs (21.8% faster)


def test_function_with_no_parameters_and_args_passed():
    def f1():
        return 1

    def f2(a):
        return a

    codeflash_output = select_matching_signature(
        [f1, f2], 100
    )  # 4.35μs -> 3.60μs (20.8% faster)


# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.